"""
    实现一个两层的神经网络
"""

import os ,sys

from common.loss_functions import cross_entropy_error
from dataset.mnist import load_mnist

sys.path.append(os.pardir)

from common.functions import *
from common.numerial_gradient import numerical_gradient

class TwoLayerNet:
    def __init__(self, input_size, hidden_size, output_size,weight_init_std=0.01):
        # 初始化权重 W
        self.params = {}
        self.params['W1'] = weight_init_std * np.random.randn(input_size, hidden_size)
        self.params['b1'] = np.zeros(hidden_size)
        self.params['W2'] = weight_init_std * np.random.randn(hidden_size, output_size)
        self.params['b2'] = np.zeros(output_size)

    def predict(self, x):
        W1,W2 = self.params['W1'],self.params['W2']
        b1,b2 = self.params['b1'],self.params['b2']

        a1 = np.dot(x, W1) + b1
        z1 = sigmoid_function(a1)

        a2 = np.dot(z1,W2) + b2
        z2 = softmax_function_v2(a2)

        return z2

    def loss(self,x,t):
        """
        损失
        :param x:   输入值
        :param t:   监督数据
        :return:
        """
        y = self.predict(x)
        # 交叉熵损失
        return cross_entropy_error(y,t)

    def accuracy(self,x,t):
        y = self.predict(x)
        y = np.argmax(y, axis=1)
        t = np.argmax(t, axis=1)

        # 准确率
        accuracy = (sum(y == t) / float(x.shape[0]))

        return accuracy

    def numerical_gradient(self,x,t):
        loss_w = lambda W : self.loss(x,t)

        grads = {}
        grads['W1'] = numerical_gradient(loss_w,self.params['W1'])
        grads['W2'] = numerical_gradient(loss_w,self.params['W2'])
        grads['b1'] = numerical_gradient(loss_w,self.params['b1'])
        grads['b2'] = numerical_gradient(loss_w,self.params['b2'])

        return grads


def test_mini_batch():

    (x_train,t_train),(x_test,t_test) = load_mnist(normalize=True,one_hot_label=True)

    train_loss_list = []

    # 超参数的设置
    # iters_num = 10000
    iters_num = 10
    train_size = x_train.shape[0]
    batch_size = 100
    learning_rate = 0.1
    network = TwoLayerNet(784,50,10)

    for i in range(iters_num):
        # 获取一批数据
        # 这里拿到的是索引值
        batch_mask = np.random.choice(train_size, batch_size)
        # 获取数据
        x_batch = x_train[batch_mask]
        t_batch = t_train[batch_mask]

        # 计算梯度
        grad = network.numerical_gradient(x_batch,t_batch)

        for key in ('W1', 'b1', 'W2', 'b2'):
            network.params[key] = network.params[key] - learning_rate * grad[key]

        loss = network.loss(x_batch,t_batch)
        train_loss_list.append(loss)
        print(loss)

if __name__ == '__main__':
    test_mini_batch()